AI Learning Series · Part 22

The KV-Cache Design Space

MHA, MQA, GQA, MLA & friends. Doc 06 derived the KV cache — memoization of history — and its memory tax; doc 08 showed the bandwidth wall it lives behind. This is the dedicated, mathematical taxonomy of every way engineers have invented to make that cache smaller.

Optimization Stack
→
KV Cache Types
→
Harness Architectures

01 The Big Picture

Doc 10 sketched the optimization stack and name-dropped GQA and MLA in one paragraph each. Here is the full design space, with the arithmetic that separates them.

Everything in inference loops back to one tensor: the KV cache. It is the only model state that grows with your conversation, the only state that must be read every decode step (doc 08's wall), and the reason a 32K-context conversation can cost more memory than the 70-billion weights underneath it. The weights are fixed; the cache is a design variable — and an entire subfield of architecture research is the question: what is the smallest tensor we can store per token such that attention can still reconstruct the keys and values it needs?

The answers form a spectrum: store everything (MHA), share copies (MQA, GQA), compress and regenerate (MLA), share across layers (YOCO), store less history (sliding windows), waste less of what you store (PagedAttention), or store fewer bits per element (KV quantization). Each point on the spectrum trades a different currency: memory, bandwidth, quality, or serving complexity.

02 What the Cache Is — a Tensor, Not a Dict

The KV cache is not a key-value store in the programmer's sense. It is a dense 4-D tensor with shape [L, n_kv, s, d_head] per K and per V — a fixed-size slab of VRAM addressed by (layer, head, position, dimension). "Key" and "value" mean attention keys and values, not hash-map keys and values.

Three consequences follow from it being a tensor:

Its size is chosen at training time. The number of KV heads n_kv, the head dimension d_head, and the layer count L are baked into the weights. You cannot convert a trained MHA model to GQA at serve time — the query heads have learned to expect their private keys. (MQA "uptraining" exists but is a retraining, not a switch.)
It is read sequentially and fully. Every decode step streams the whole slab — all L layers, all positions, all heads — through HBM exactly once per token generated (doc 06, doc 08). Cache bytes → decode latency, directly.
Its layout is a performance contract. K is stored transposed — per head, shape [d_head, s] instead of [s, d_head] — so the score computation S = Q·Kᵀ is a clean GEMM over contiguous memory, and V stays [s, d_head] for the value-weighted sum. Contiguous runs along the s axis are what let FlashAttention-style kernels read whole histories with coalesced, high-bandwidth transactions. Every trick in this doc — paging, quantization, eviction — must preserve or explicitly pay for that contiguity.
📐
The design question, formally: per token, per layer, MHA stores 2 · n_h · d_head elements. Every design below is an answer to "can we store fewer than that?" — and at what cost to the attention output.

03 Why — The Wall, Times Three

Doc 08 established that decode is memory-bandwidth-bound: the GPU reads weights and the KV cache from HBM for every single token. Cache design matters because three multipliers compound:

// the serving squeeze concurrency = VRAM_free ÷ KV_bytes_per_seq // how many users fit at all tok/s/user ≈ BW ÷ (model_bytes + KV_bytes_per_seq) // how fast each one decodes cost/token ∝ GPU_seconds ÷ tokens_served // what you're charged

Context length × batch size × cache size ⇒ throughput. Shrink the per-token cache and every factor improves at once: more sequences fit in VRAM (bigger effective batch), each decode step reads fewer bytes (faster tokens), and the freed capacity can hold other users' caches — which is precisely what makes prompt-caching economics possible (doc 07). A provider whose cache is 8× smaller can keep 8× more prefixes resident, raising cache-hit rates and deepening the cached-token discount you see on the price card.

That is the "why": the KV cache is the scarce resource of serving, and the designs in this doc are its compression algorithms.

04 How — One Request, Three Designs

Watch the same 8-token request write its cache under MHA, GQA, and MLA, one decoder layer at a time.

One request · 8 tokens · one decoder layer · slots = stored K/V vectors t₁ t₂ t₃ t₄ t₅ t₆ t₇ t₈ the request MHA — n_h = 64 KV heads each cell = 4 KV heads · 64 written per token · 2·64·128·2 B = 32 KB / token / layer GQA — n_kv = 8 never written 8 KV heads, each shared by 8 query heads · n_kv/n_h = 8/64 → 4 KB / token / layer (÷8) MLA — one latent per token c_t ∈ ℝ⁵¹² — latent that carries BOTH K and V kᵣₒₚₑ ∈ ℝ⁶⁴ (512 + 64) elems · 2 B ≈ 1.1 KB / token / layer — K and V compressed jointly ⚠ restore: k = W_UP · c per position — a small matmul on every read (compute buys memory) same request, same layer: 64 slots → 8 slots → 1 latent + rope · 1× → ⅛× → ~1/28× the bytes and every stored byte is re-read from HBM on each decode step (doc 08) — so this ratio is also the speed ratio

The three rows are the spine of this doc; everything that follows either interpolates between them (MQA, GQA's knob), pushes further (cross-layer sharing), or changes how the bytes are stored rather than how many (paging, quantization, eviction).

05 The Taxonomy

MHA — the baseline

Multi-Head Attention. Every query head owns a private K and V. Full representational capacity — each head learns its own "address book" — at the highest memory cost: bytes/token = 2 · L · n_h · d_head · s · b.

✓ Pros

Maximum quality headroom; simplest kernels; no restore compute; every attention implementation supports it.

✗ Cons

Biggest cache — the 1.0× every other design is measured against. At long context, memory and decode bandwidth dominate everything.

Where used: original Transformer, GPT-2, BERT-era encoders, Llama-1 — and modern small models where the cache is small enough not to matter.

MQA — one shared address book

Multi-Query Attention (Shazeer, 2019): all query heads read from a single shared K/V head. The cache shrinks exactly ÷n_h — 64× for a 64-head model — because n_kv = 1.

✓ Pros

Maximum savings; also shrinks the K/V weight matrices and makes decode GEMMs fatter (better GPU utilization).

✗ Cons

Measurable quality drop — all queries share one view of history, costing diversity exactly where attention needs it.

Where used: PaLM, Falcon, StarCoder.

GQA — the interpolation knob

Grouped-Query Attention: partition the n_h query heads into groups; each group shares one K/V head. n_kv is a continuous dial between MQA (n_kv = 1) and MHA (n_kv = n_h), with the memory ratio simply n_kv / n_h. Llama-2/3 70B: 8 KV heads serving 64 query heads → 4 KB/token/layer instead of 32 KB, at near-MHA quality. This is the industry's default compromise.

MLA — compress, then regenerate

Multi-head Latent Attention (DeepSeek-V2/V3): instead of sharing whole heads, compress all K and V jointly into one low-rank latent vector per token. See the derivation in §06. ~28× smaller than MHA at comparable width, at near-lossless quality — the price is per-position restore compute and custom kernels.

Cross-layer / YOCO-style shared KV

"You Only Cache Once." Stack a few shared-KV attention layers whose cache is reused by many downstream layers: the layer-count factor in the sizing formula partially collapses. The extreme version caches once for the whole model. Trade: layers no longer build progressively richer keys, which constrains architecture design — but the savings multiply with everything else.

Sliding-window + global hybrid layers

Mistral, Gemma: most layers attend only to the last w tokens (4K–8K), a few layers attend globally. Local layers cap the cache at w per layer and make the attention matrix O(s·w) instead of O(s²); the global layers — far fewer — preserve long-range recall. Per-layer storage stops growing with s for the majority of the model.

PagedAttention — waste nothing (vLLM)

Pre-allocation forces reserving s_max per sequence; real sequences rarely fill it, and the fragmentation + reservation waste historically discarded 60–80% of KV memory. PagedAttention stores the cache in fixed-size blocks (classic block size: 16 tokens) with a page table — OS virtual memory for tensors:

// logical position → physical block logical_block = pos ÷ block_size physical_block = block_table[seq_id][logical_block] slot = physical_block × block_size + (pos mod block_size)

It stores the same bytes — it's the design that makes the other designs' savings actually bankable, and its copy-on-write blocks are what make prefix sharing (doc 07) and parallel sampling cheap.

KV quantization, eviction, offload

FP8 / INT4 KV: halve or quarter b. The subtlety: V tolerates coarse quantization, but K does not — softmax divides by the temperature and is acutely sensitive to outliers in specific K channels, so K needs per-channel scales (one scale per head-dimension) while V gets away with per-token scales. Eviction (H₂O, SNAP): drop tokens with low accumulated attention mass, keeping "heavy hitters" + recency — the cache stops being a lossless record. Prefix caching + NVMe offload: move cold prefixes down GPU HBM → CPU DRAM → NVMe, keyed by token-prefix hash — doc 07's mechanism, enabled by small caches.

06 The Math

One formula to size them all

// unified KV-cache sizing — every design in this doc is a choice of (d_kv, n_kv, L_eff, b) KV_bytes(request) = 2 · L · s · d_kv · n_kv · b L = layers storing the cache s = sequence length d_kv = per-head key/value dim b = bytes per element (2 = fp16, 1 = fp8) n_kv = KV heads (n_h for MHA, groups for GQA, 1 for MQA) the leading 2 = K and V — MLA replaces this row entirely (single latent, no ×2)

Worked table — one 70B-class sequence, s = 32K, fp16

Model shape: L = 80, d_head = 128, n_h = 64 (so d_model = 8192). Plug in:

Designn_kv / latentElems / token / layerBytes / token / layerGB @ 32K ctxvs MHA
MHAn_kv = 642·64·128 = 16,38432 KB2·80·32768·128·64·2 ≈ 85.9 GB1×
GQA (8 groups)n_kv = 82·8·128 = 2,0484 KB≈ 10.7 GB⅛×
MQAn_kv = 12·1·128 = 2560.5 KB≈ 1.34 GB1/64×
MLAd_c + d_rope = 576576≈1.1 KB80·32768·576·2 ≈ 3.0 GB~1/28×

Read the last two rows together: MLA's cache is bigger than MQA's per token, yet delivers near-MHA quality — it pays 576 elements to keep a compressed copy of all heads' information instead of 256 elements that literally share one head. That is the recurring theme: the knob isn't "how few bytes," it's "how much information survives the compression."

MLA derivation sketch

MHA keeps, per token t, n_h separate k-vectors and v-vectors — 2·n_h·d_h numbers, most of them redundant across heads. MLA's bet: the joint K/V content of a token lives in a much lower-dimensional subspace. So project it down once, at write time:

// write path — one matrix multiply, stored forever c_t = W_DK · x_t // down-projection: d_model → d_c ≈ 512 (the stored latent) kᵣₒₚₑ = W_KR · x_t // decoupled RoPE key: d_model → d_rope = 64 (stored separately) // restore path — on demand, per position k_t = W_UK · c_t // up-project back to n_h heads' keys v_t = W_UV · c_t // ...and values // why the decoupled rope key exists: RoPE rotates k_t position-wise, // so a rotated key can NOT be a linear function of the (unrotated) latent. // Rope is carried in a small separate 64-d slice, shared across heads.

The elegance is in the absorption trick: attention scores need qᵀ·k_t = qᵀ·W_UK·c_t = (W_UKᵀ·q)ᵀ·c_t. Fold W_UK into the query projection at kernel time and you never materialize k_t at all — you score directly against the 512-d latent. What remains is one small matmul per position per layer on the restore path: compute and kernel complexity bought with a ~28× memory discount. The capacity check is one line:

// latent vs full capacity, same d_model = 8192 n_stored = d_c + d_rope = 512 + 64 = 576 ≪ n_stored(MHA) = 2 · n_h · d_h = 16,384 576 / 16384 ≈ 1/28 // the row in the table above
🧮
Compose the levers: the formula multiplies. GQA (⅛) × FP8 cache (½) × sliding window (w/s) can put a 32K context's cache under 1 GB — which is what lets vLLM pack hundreds of sessions, and what makes a provider's cached-token discount (doc 07) economically sustainable.

07 Engineering Takeaways — Which to Pick When

SituationPickWhy
Short context, quality-critical, small modelMHACache is small enough that capacity buys real quality
General-purpose long-context servingGQA, n_kv ≈ n_h/8The industry equilibrium: ~⅛ memory, ~no quality loss
Extreme batch / code models, quality slack existsMQA1/n_kv = 1 savings; Falcon/StarCoder accepted the trade
Very long context at fleet scaleMLA~28× + quality — if you can pay the kernel-complexity tax
Cache doesn't fit VRAM at allPaging + quantization + offloadOrthogonal to the above — multiplies with every row
These are training-time decisions. Reading a model card — "GQA 8 KV heads", "MLA, d_c = 512" — tells you its cache economics before you send a token, exactly like reading "MoE, 37B active" told you its FLOPs economics in doc 10 and doc 21.
Serving complexity is a real cost column. MHA and GQA run on stock kernels; MLA needs absorbed-attention kernels, quantized caches need matching dequant kernels, paged caches need a block-table-aware allocator. The design you pick is also the engineering burden you pick.
Cache design is why cache pricing exists. A cache-hit discount is the provider sharing the savings of a small, cheap, long-resident working set. GQA/MLA made the working set small; PagedAttention made residency fragment-free; prefix hashing (doc 07) made reuse possible. The ~0.1× cached-input price line is this entire doc, expressed as money.
Decode speed follows the same ratio. Every stored byte is read every step (doc 08): a ⅛× cache is roughly an 8× lighter decode read. When you benchmark providers, cache design explains token-speed differences that model size alone doesn't. Where this slot-filling fits into prefill vs decode strategy: doc 24.

08 Mental Models

Main memory of attention, register-file compaction

The KV cache is attention's main memory: every query must fetch from it before it can think. KV-cache design is then register-file compaction — the same game RISC architects play of keeping the hottest state in the fewest bytes, because every access is on the critical path. Lets you reason about: why cache bytes convert 1:1 into decode latency, and why compression here beats almost any other optimization.

A register file is fixed-size; this memory grows with conversation length — the pressure compounds instead of staying constant.
MLA is a ZIP archive for the address book

MHA stores the full address book (all heads' K,V). MLA stores a compressed archive (the latent) plus instructions to regenerate each page on demand — paying a small unzip cost (W_UK·c) each time a page is read. Lets you reason about: the memory-vs-compute trade, and why restore cost shows up in kernel engineering, not in model quality.

ZIP is lossless; the latent is lossy by design — the bet is that attention can't tell the difference.
GQA is carpooling

64 query heads commuting to the same history. MHA gives each its own car (private K/V); MQA puts everyone in one bus (one shared head, crowded and slow-witted); GQA runs 8 carpools — nearly as cheap as the bus, nearly as comfortable as private cars. Lets you reason about: n_kv as a continuous dial, not a binary choice.

Carpool routes are fixed at training time — you can't re-pool a trained model's heads at serve time.

09 Common Misconceptions

"GQA hurts quality as much as MQA." The quality cliff is near MQA (n_kv = 1); at n_kv ≈ 8 the loss is small enough that it became the default for every major open model. The dial's ends are very different from its middle.

"MLA is just fancy GQA." No. GQA shares whole heads but keeps every dimension of the heads it keeps. MLA compresses across the dimension too, jointly for K and V, into one latent — a fundamentally different (low-rank, restorable) representation, with the restore matmul and decoupled-RoPE machinery GQA doesn't have.

"KV quantization is like weight quantization." Weights are static and forgiving. K vectors feed a softmax that is sharply sensitive to outliers in specific channels — hence per-channel scales for K, per-token for V, and quality degradation that appears only at long range, exactly where you wanted the savings.

"PagedAttention makes the cache smaller." It stores the identical bytes; it stops wasting 60–80% of the allocation on fragmentation and over-reservation. The savings are real but come from accounting, not compression — it composes with, rather than replaces, GQA/MLA/quantization.

"Sliding window means the model can't remember anything long-range." In hybrid designs (Mistral, Gemma), only most layers are local; dedicated global layers still attend across the whole context. You lose long-range per-local-layer, not long-range capability overall — though retrieval-heavy tasks do measurably degrade versus full attention.

🧭
The dot this doc connects: doc 06's memoization tax → doc 08's bandwidth wall → the compression designs here → doc 07's cached-token prices. The whole stack — model design, hardware limits, and the line on your invoice — is one causal chain. Next: doc 24 shows how prefill and decode strategies spend this newly-affordable cache budget, and harness architectures build on top of it.